#!/usr/bin/env python3
"""
通用评估模块 - CLIP 系列模型的检索评估函数
==================================================

本模块提供了一套完整的、标准化的图像-文本检索评估工具，专为 CLIP 系列模型设计。
支持各种 CLIP 变体模型（如 CN-CLIP, OpenAI CLIP, Chinese-CLIP 等）的性能评估。

核心功能包括：
1. 双向检索评估：Image-to-Text (I2T) 和 Text-to-Image (T2I)
2. 多种评估指标：Recall@1, Recall@5, Recall@10
3. 灵活的参数配置：温度缩放、特征归一化、自定义 K 值
4. 结构化数据输出：便于日志记录和结果分析
5. 高性能实现：基于 PyTorch 张量操作，支持 GPU 加速

理论背景：
----------
在 CLIP 模型的评估中，我们通常关注两个核心任务：
- 图像检索任务：给定文本描述，找到最相关的图像
- 文本检索任务：给定图像，找到最相关的文本描述

评估指标 Recall@K 定义：
- 在 top-K 检索结果中包含正确答案的查询比例
- K=1 表示第一名就是正确答案的比例（最严格）
- K=5, K=10 表示前5名、前10名中包含正确答案的比例

数学原理：
----------
对于批次大小为 N 的样本：
1. 计算相似度矩阵 S ∈ R^(N×N)，其中 S[i,j] = cosine_similarity(image_i, text_j)
2. 对角线元素 S[i,i] 表示第 i 个图像与第 i 个文本的相似度（正确配对）
3. I2T Recall@K = (1/N) * Σ_i I(rank(S[i,i] in S[i,:]) ≤ K)
4. T2I Recall@K = (1/N) * Σ_j I(rank(S[j,j] in S[:,j]) ≤ K)

其中 I(·) 是指示函数，rank(·) 表示排名（1表示最高相似度）。

使用场景：
----------
- 模型训练过程中的验证评估
- 不同模型或配置的性能对比
- 研究论文中的标准化评估报告
- 工业应用中的模型质量监控

Author: WenwuClip Project Team
Version: 1.0.0
License: MIT
"""

from typing import Dict, List, Tuple, Union, Optional
import torch
import torch.nn.functional as F
from dataclasses import dataclass

# ==========================================
# 数据结构定义
# ==========================================


@dataclass
class RetrievalMetrics:
    """
    检索评估指标数据类
    
    这个数据类存储了所有的 CLIP 检索评估指标，提供了结构化的数据访问接口。
    包含两个检索方向的指标以及它们的平均值，便于全面评估模型性能。
    
    属性说明：
    ----------
    I2T (Image-to-Text) 指标：
        i2t_r1:  图像检索文本的 Recall@1 - 第一名就是正确答案的比例
        i2t_r5:  图像检索文本的 Recall@5 - 前5名包含正确答案的比例  
        i2t_r10: 图像检索文本的 Recall@10 - 前10名包含正确答案的比例
        
    T2I (Text-to-Image) 指标：
        t2i_r1:  文本检索图像的 Recall@1 - 第一名就是正确答案的比例
        t2i_r5:  文本检索图像的 Recall@5 - 前5名包含正确答案的比例
        t2i_r10: 文本检索图像的 Recall@10 - 前10名包含正确答案的比例
        
    平均指标（最常用的综合性能指标）：
        mean_r1:  I2T 和 T2I 的 R@1 平均值 - 最重要的单一性能指标
        mean_r5:  I2T 和 T2I 的 R@5 平均值  
        mean_r10: I2T 和 T2I 的 R@10 平均值
    
    使用示例：
    ----------
    >>> metrics = RetrievalMetrics(...)
    >>> print(f"模型整体性能: Mean R@1 = {metrics.mean_r1:.3f}")
    >>> if metrics.i2t_r1 > metrics.t2i_r1:
    ...     print("图像检索文本的能力更强")
    >>> else:
    ...     print("文本检索图像的能力更强")
    """
    
    # ====== Image-to-Text (I2T) 检索指标 ======
    # 给定图像，检索最相关的文本描述
    i2t_r1: float   # Recall@1:  最严格的指标，要求第一名就是正确答案
    i2t_r5: float   # Recall@5:  前5名中包含正确答案即可
    i2t_r10: float  # Recall@10: 前10名中包含正确答案即可
    
    # ====== Text-to-Image (T2I) 检索指标 ======  
    # 给定文本描述，检索最相关的图像
    t2i_r1: float   # Recall@1:  最严格的指标，要求第一名就是正确答案
    t2i_r5: float   # Recall@5:  前5名中包含正确答案即可
    t2i_r10: float  # Recall@10: 前10名中包含正确答案即可
    
    # ====== 综合平均指标 ======
    # 这些是最常报告的综合性能指标，结合了两个检索方向的性能
    mean_r1: float   # 平均 R@1:  最重要的单一指标，论文中经常用于模型比较
    mean_r5: float   # 平均 R@5:  较为宽松的综合指标
    mean_r10: float  # 平均 R@10: 最宽松的综合指标
    
    def to_dict(self) -> Dict[str, float]:
        """
        转换为字典格式，便于日志记录和数据序列化
        
        这个方法将结构化的指标数据转换为扁平的字典格式，
        便于与日志系统（如 TensorBoard, WandB）或数据库进行集成。
        
        Returns:
        --------
        Dict[str, float]: 包含所有指标的字典，键名与属性名一致
        
        使用示例:
        ---------
        >>> metrics = RetrievalMetrics(i2t_r1=0.5, ...)
        >>> metrics_dict = metrics.to_dict()
        >>> logger.log(metrics_dict)  # 直接用于日志记录
        >>> json.dump(metrics_dict, f)  # 保存为 JSON 文件
        """
        return {
            # Image-to-Text 指标
            'i2t_r1': self.i2t_r1,      # 图像→文本 Recall@1
            'i2t_r5': self.i2t_r5,      # 图像→文本 Recall@5 
            'i2t_r10': self.i2t_r10,    # 图像→文本 Recall@10
            
            # Text-to-Image 指标
            't2i_r1': self.t2i_r1,      # 文本→图像 Recall@1
            't2i_r5': self.t2i_r5,      # 文本→图像 Recall@5
            't2i_r10': self.t2i_r10,    # 文本→图像 Recall@10
            
            # 平均指标（最重要的综合指标）
            'mean_r1': self.mean_r1,    # 平均 Recall@1 (最常用指标)
            'mean_r5': self.mean_r5,    # 平均 Recall@5
            'mean_r10': self.mean_r10,  # 平均 Recall@10
        }


# ==========================================
# 核心评估函数
# ==========================================

def compute_clip_retrieval_metrics(
    image_features: torch.Tensor,
    text_features: torch.Tensor,
    k_values: List[int] = [1, 5, 10],
    temperature: float = 0.07,
    normalize: bool = True
) -> RetrievalMetrics:
    """
    计算 CLIP 模型的双向检索评估指标 - 核心评估函数
    
    这是本模块的核心函数，实现了完整的 CLIP 模型检索性能评估流程。
    支持批量处理、GPU 加速，并提供了灵活的参数配置选项。
    
    算法流程详解：
    ==============
    1. 输入验证：检查张量维度匹配性
    2. 特征预处理：可选的 L2 归一化 
    3. 相似度计算：计算 N×N 相似度矩阵
    4. 双向排序：分别对行（I2T）和列（T2I）进行排序
    5. 召回率统计：统计正确答案在 top-K 中的比例
    6. 结果汇总：计算各种指标并返回结构化结果
    
    理论基础：
    ==========
    CLIP 模型通过对比学习将图像和文本映射到同一个语义空间中，
    理想情况下，配对的图像和文本应该在这个空间中距离最近。
    
    相似度矩阵的含义：
    - S[i,j] = cosine_similarity(image_i, text_j)
    - 对角线 S[i,i] 是第 i 个图像与第 i 个文本的相似度（正确配对）
    - 非对角线元素是不正确的配对相似度（负样本）
    
    Recall@K 的计算：
    - 对每个查询，将所有候选按相似度排序
    - 检查正确答案的排名是否 ≤ K
    - 计算所有查询中满足条件的比例
    
    参数详解：
    ==========
    image_features : torch.Tensor, shape [N, D]
        图像特征张量，通常来自 CLIP 视觉编码器的输出
        - N: 批次大小（样本数量）
        - D: 特征维度（如 512, 768, 1024 等）
        - 数据类型: torch.float32 或 torch.float16
        - 设备: 支持 CPU 和 GPU（会自动适配输入张量的设备）
        
    text_features : torch.Tensor, shape [N, D]  
        文本特征张量，通常来自 CLIP 文本编码器的输出
        - 必须与 image_features 具有相同的批次大小和特征维度
        - 第 i 个文本特征应该与第 i 个图像特征构成正确配对
        
    k_values : List[int], default [1, 5, 10]
        要计算的 K 值列表，指定 top-K 检索的 K 值
        - 常用值：[1, 5, 10] 对应严格、中等、宽松的评估标准
        - 如果 K > N（批次大小），会自动设置 Recall@K = 1.0
        - 支持任意正整数，但建议不超过批次大小
        
    temperature : float, default 0.07
        温度参数，用于缩放相似度分数。
        - **重要**: 此参数与模型训练时的 `logit_scale` 紧密相关。
        - 关系: `temperature` = 1 / `logit_scale.exp()`。
        - `0.07` 是 OpenAI CLIP 原始论文中的标准值，对应 `logit_scale` ≈ 2.6592
        - 错误地设置此参数将导致评估结果不准确。
        - 数学上：scaled_similarity = similarity / temperature
        
    normalize : bool, default True
        是否对特征进行 L2 归一化
        - True（推荐）: 计算余弦相似度，范围在 [-1, 1]
        - False: 计算点积相似度，范围依赖于特征的模长
        - 大多数 CLIP 模型都需要归一化来获得最佳性能
    
    返回值详解：
    ============
    RetrievalMetrics 对象，包含以下指标：
    
    双向检索指标：
        - i2t_r1, i2t_r5, i2t_r10: Image-to-Text 的 Recall@K
        - t2i_r1, t2i_r5, t2i_r10: Text-to-Image 的 Recall@K
        - mean_r1, mean_r5, mean_r10: 双向平均的 Recall@K
    
    指标解释：
        - Recall@1: 最严格，第一名必须是正确答案
        - Recall@5: 前5名中包含正确答案即可
        - Recall@10: 前10名中包含正确答案即可
        - Mean: 平均值是最常用的综合性能指标
    
    异常处理：
    ==========
    ValueError: 
        - 当 image_features.shape[0] ≠ text_features.shape[0] 时（批次大小不匹配）
        - 当 image_features.shape[1] ≠ text_features.shape[1] 时（特征维度不匹配）
    
    性能考虑：
    ==========
    - 时间复杂度：O(N²) 用于相似度矩阵计算，O(N²log(N)) 用于排序
    - 空间复杂度：O(N²) 用于存储相似度矩阵
    - GPU 加速：完全支持，自动使用输入张量的设备
    - 内存优化：对于大批次，考虑分块处理以避免内存不足
    
    使用示例：
    ==========
    # 基础用法
    >>> image_feat = torch.randn(32, 512)  # 32个图像，512维特征
    >>> text_feat = torch.randn(32, 512)   # 32个文本，512维特征
    >>> metrics = compute_clip_retrieval_metrics(image_feat, text_feat)
    >>> print(f"平均 R@1: {metrics.mean_r1:.3f}")
    
    # 自定义参数
    >>> metrics = compute_clip_retrieval_metrics(
    ...     image_feat, text_feat,
    ...     k_values=[1, 3, 5, 10],      # 自定义 K 值
    ...     temperature=0.1,             # 更尖锐的分布
    ...     normalize=True               # 使用余弦相似度
    ... )
    
    # GPU 使用
    >>> if torch.cuda.is_available():
    ...     image_feat = image_feat.cuda()
    ...     text_feat = text_feat.cuda()
    ...     metrics = compute_clip_retrieval_metrics(image_feat, text_feat)
    
    # 与其他工具集成
    >>> metrics_dict = metrics.to_dict()
    >>> wandb.log(metrics_dict)  # 记录到 Weights & Biases
    >>> tensorboard.add_scalars('eval_metrics', metrics_dict, step)
    
    性能基准参考：
    ==============
    不同模型在标准数据集上的典型性能（仅供参考）：
    - 随机基线：Mean R@1 ≈ 1/N（N为候选数量）
    - 弱 CLIP 模型：Mean R@1 ≈ 0.1-0.3
    - 中等 CLIP 模型：Mean R@1 ≈ 0.3-0.6  
    - 强 CLIP 模型：Mean R@1 ≈ 0.6-0.8
    - SOTA 模型：Mean R@1 ≈ 0.8+
    
    注意事项：
    ==========
    1. 输入的图像和文本特征必须是配对的（第i个图像对应第i个文本）
    2. 特征应该来自同一个模型的编码器输出
    3. 对于不同模型，可能需要调整 temperature 参数
    4. 批次大小会影响评估的难度（更大的批次意味着更多负样本）
    5. 该函数假设每个图像只有一个正确的文本描述（1对1配对）
    """
    
    # ==========================================
    # 第一步：输入验证和数据预处理
    # ==========================================
    
    # 验证输入张量的批次大小是否一致
    # 这是最基本的要求，因为我们需要假设第i个图像与第i个文本配对
    if image_features.shape[0] != text_features.shape[0]:
        raise ValueError(
            f"Batch size mismatch: images={image_features.shape[0]}, texts={text_features.shape[0]}. "
            f"Image and text features must have the same batch size for paired evaluation."
        )
    
    # 验证特征维度是否一致
    # 图像和文本特征必须在同一个语义空间中，因此维度必须相同
    if image_features.shape[1] != text_features.shape[1]:
        raise ValueError(
            f"Feature dim mismatch: images={image_features.shape[1]}, texts={text_features.shape[1]}. "
            f"Image and text features must have the same feature dimension."
        )
    
    # 提取关键信息
    n = image_features.shape[0]  # 批次大小，即样本数量
    device = image_features.device  # 自动检测设备（CPU 或 GPU）
    
    # ==========================================
    # 第二步：特征归一化（可选但推荐）
    # ==========================================
    
    if normalize:
        """
        L2 归一化的作用和重要性：
        1. 将特征向量的模长归一化为1，消除模长差异的影响
        2. 使得相似度计算变为余弦相似度：cos(θ) = (a·b) / (||a||·||b||)
        3. 余弦相似度范围在 [-1, 1]，更稳定和可解释
        4. 大多数预训练的 CLIP 模型都是在归一化特征上训练的
        
        数学公式：normalized_x = x / ||x||_2
        其中 ||x||_2 = sqrt(Σ(x_i²)) 是 L2 范数
        """
        # 对图像特征进行 L2 归一化，p=2 表示 L2 范数，dim=1 表示对特征维度归一化
        image_features = F.normalize(image_features, p=2, dim=1)
        # 对文本特征进行相同的归一化处理
        text_features = F.normalize(text_features, p=2, dim=1)
    
    # ==========================================
    # 第三步：计算相似度矩阵
    # ==========================================
    
    """
    相似度矩阵是整个评估流程的核心数据结构：
    
    矩阵形状: [N, N]
    矩阵含义: similarities[i,j] = similarity(image_i, text_j)
    
    对角线元素 similarities[i,i]：
    - 表示第 i 个图像与第 i 个文本的相似度
    - 这是正确配对的相似度，理想情况下应该是该行/列的最大值
    
    非对角线元素 similarities[i,j] (i≠j)：
    - 表示不正确配对的相似度
    - 这些是负样本，用于测试模型是否能区分正确和错误的配对
    
    温度参数的作用：
    - temperature > 1: 使相似度分布更平滑，减小差异
    - temperature < 1: 使相似度分布更尖锐，放大差异
    - temperature = 1: 不改变原始相似度
    
    实际计算过程：
    1. torch.matmul(image_features, text_features.t()) 计算矩阵乘积
    2. text_features.t() 将文本特征转置为 [D, N]
    3. 结果是 [N, N] 的相似度矩阵
    4. 除以 temperature 进行温度缩放
    """
    similarities = torch.matmul(image_features, text_features.t()) / temperature
    
    # 相似度矩阵的形状检查（调试时有用）
    assert similarities.shape == (n, n), f"相似度矩阵形状错误: {similarities.shape}, 期望: ({n}, {n})"
    
    # ==========================================
    # 第四步：双向检索评估
    # ==========================================
    
    # 初始化指标字典，用于存储所有计算出的召回率指标
    metrics_dict = {}
    
    # 遍历所有要计算的 K 值（通常是 [1, 5, 10]）
    for k in k_values:
        
        # ==========================================
        # 边界条件处理：K > N 的情况
        # ==========================================
        
        if k > n:
            """
            当 K 大于批次大小 N 时的特殊处理：
            
            在这种情况下，所有候选都会被包含在 top-K 中，
            因此召回率必然是 100%（每个查询都能找到正确答案）。
            
            这种情况在小批次测试或 K 值设置过大时会出现。
            例如：批次大小为 8，但要计算 Recall@10。
            """
            metrics_dict[f'i2t_r{k}'] = 1.0  # Image-to-Text 完美召回
            metrics_dict[f't2i_r{k}'] = 1.0  # Text-to-Image 完美召回
            continue  # 跳过实际计算，直接处理下一个 K 值
            
        # ==========================================
        # Image-to-Text (I2T) 检索评估
        # ==========================================
        
        """
        I2T 检索的核心思想：
        
        给定一个图像，在所有 N 个文本中找到最相似的前 K 个。
        然后检查正确的文本（与该图像配对的文本）是否在这 K 个候选中。
        
        具体步骤：
        1. 对相似度矩阵的每一行进行排序（每行代表一个图像与所有文本的相似度）
        2. 取每行的前 K 个最大值的索引（即最相似的 K 个文本的索引）
        3. 检查每个图像的正确文本索引是否在其 top-K 列表中
        4. 计算成功找到正确文本的图像比例
        
        数学表示：
        对于第 i 个图像，其正确文本索引为 i
        I2T_success_i = 1 if i ∈ top_K_texts_for_image_i else 0
        I2T_Recall@K = (1/N) * Σ(I2T_success_i)
        """
        
        # 步骤1：对每一行进行排序，获取每个图像最相似的 K 个文本的索引
        # torch.argsort(similarities, dim=1, descending=True) 的详细解释：
        # - similarities: [N, N] 相似度矩阵
        # - dim=1: 对第1维（列维度）进行排序，即对每行独立排序
        # - descending=True: 降序排列，最相似的排在前面
        # - 返回值: [N, N] 排序后的索引矩阵
        # - [:, :k]: 只取每行的前 K 个索引，结果形状 [N, K]
        i2t_ranked_indices = torch.argsort(similarities, dim=1, descending=True)[:, :k]
        
        # 步骤2：构造目标索引（正确答案）
        # 对于第 i 个图像，其正确的文本索引就是 i（对角线位置）
        # torch.arange(n, device=device): 创建 [0, 1, 2, ..., n-1] 张量
        # .unsqueeze(1): 将形状从 [N] 改为 [N, 1]，便于后续广播比较
        targets_i2t = torch.arange(n, device=device).unsqueeze(1)  # 形状: [N, 1]
        
        # 步骤3：检查正确答案是否在 top-K 候选中
        # i2t_ranked_indices: [N, K] - 每个图像的 top-K 文本索引
        # targets_i2t: [N, 1] - 每个图像的正确文本索引
        # 广播比较: [N, K] == [N, 1] -> [N, K] 布尔矩阵
        # torch.any(..., dim=1): 对每行进行 OR 操作，如果该行有任何 True，则该图像检索成功
        # 结果: [N] 布尔张量，表示每个图像是否成功检索到正确文本
        i2t_correct = torch.any(i2t_ranked_indices == targets_i2t, dim=1)  # 形状: [N]
        
        # 步骤4：计算召回率
        # .float(): 将布尔值转为浮点数 (True -> 1.0, False -> 0.0)
        # .mean(): 计算平均值，即成功的比例
        # .item(): 将单元素张量转为 Python 标量
        i2t_recall = i2t_correct.float().mean().item()
        
        # ==========================================
        # Text-to-Image (T2I) 检索评估  
        # ==========================================
        
        """
        T2I 检索的核心思想：
        
        给定一个文本，在所有 N 个图像中找到最相似的前 K 个。
        然后检查正确的图像（与该文本配对的图像）是否在这 K 个候选中。
        
        与 I2T 的区别：
        - I2T 是对矩阵的行进行操作（每行是一个图像的查询）
        - T2I 是对矩阵的列进行操作（每列是一个文本的查询）
        
        具体步骤：
        1. 对相似度矩阵的每一列进行排序（每列代表所有图像与一个文本的相似度）
        2. 取每列的前 K 个最大值的索引（即最相似的 K 个图像的索引）
        3. 检查每个文本的正确图像索引是否在其 top-K 列表中
        4. 计算成功找到正确图像的文本比例
        """
        
        # 步骤1：对每一列进行排序，获取每个文本最相似的 K 个图像的索引
        # torch.argsort(similarities, dim=0, descending=True) 的详细解释：
        # - dim=0: 对第0维（行维度）进行排序，即对每列独立排序
        # - 结果: [N, N] 排序后的索引矩阵
        # - [:k, :]: 只取每列的前 K 个索引，结果形状 [K, N]
        t2i_ranked_indices = torch.argsort(similarities, dim=0, descending=True)[:k, :]
        
        # 步骤2：构造目标索引（正确答案）
        # 对于第 j 个文本，其正确的图像索引就是 j（对角线位置）
        # .unsqueeze(0): 将形状从 [N] 改为 [1, N]，便于后续广播比较
        targets_t2i = torch.arange(n, device=device).unsqueeze(0)  # 形状: [1, N]
        
        # 步骤3：检查正确答案是否在 top-K 候选中
        # t2i_ranked_indices: [K, N] - 每个文本的 top-K 图像索引
        # targets_t2i: [1, N] - 每个文本的正确图像索引
        # 广播比较: [K, N] == [1, N] -> [K, N] 布尔矩阵
        # torch.any(..., dim=0): 对每列进行 OR 操作，如果该列有任何 True，则该文本检索成功
        # 结果: [N] 布尔张量，表示每个文本是否成功检索到正确图像
        t2i_correct = torch.any(t2i_ranked_indices == targets_t2i, dim=0)  # 形状: [N]
        
        # 步骤4：计算召回率
        t2i_recall = t2i_correct.float().mean().item()
        
        # ==========================================
        # 存储当前 K 值的评估结果
        # ==========================================
        
        # 将计算出的召回率存储到指标字典中
        # 键名格式：'i2t_r{k}' 和 't2i_r{k}'，如 'i2t_r1', 't2i_r5' 等
        metrics_dict[f'i2t_r{k}'] = i2t_recall  # 存储 I2T Recall@K
        metrics_dict[f't2i_r{k}'] = t2i_recall  # 存储 T2I Recall@K
    
    # ==========================================
    # 第五步：计算平均指标（综合性能评估）
    # ==========================================
    
    """
    平均指标的重要性：
    
    单一方向的指标（I2T 或 T2I）只能反映模型在特定任务上的性能，
    而平均指标能够提供更全面的模型性能评估：
    
    1. 消除方向偏差：某些模型可能在一个方向上表现更好
    2. 简化比较：用单一数值比较不同模型的整体性能  
    3. 论文报告：学术论文通常报告 Mean R@1 作为主要指标
    4. 模型选择：训练时可以用 Mean R@1 作为早停和模型保存的依据
    
    计算公式：
    Mean R@K = (I2T R@K + T2I R@K) / 2
    """
    
    # 对于每个 K 值，计算 I2T 和 T2I 的平均召回率
    for k in k_values:
        
        if k <= n:
            # 标准情况：K 不超过批次大小，有实际计算的指标
            
            # 构造指标键名
            i2t_key = f'i2t_r{k}'      # 如 'i2t_r1', 'i2t_r5'
            t2i_key = f't2i_r{k}'      # 如 't2i_r1', 't2i_r5' 
            mean_key = f'mean_r{k}'    # 如 'mean_r1', 'mean_r5'
            
            # 检查两个方向的指标是否都已计算
            if i2t_key in metrics_dict and t2i_key in metrics_dict:
                # 计算算术平均值
                metrics_dict[mean_key] = (metrics_dict[i2t_key] + metrics_dict[t2i_key]) / 2.0
            else:
                # 异常情况：某个方向的指标缺失（理论上不应该发生）
                metrics_dict[mean_key] = 1.0  # 设为完美值作为安全默认
                
        else:
            # 边界情况：K > N，直接设置为完美召回率
            metrics_dict[f'mean_r{k}'] = 1.0
    
    # ==========================================
    # 第六步：构建结构化结果对象
    # ==========================================
    
    """
    结果对象构建的考虑：
    
    1. 默认值处理：使用 .get() 方法提供默认值 0.0，防止键缺失错误
    2. 标准化输出：确保输出对象包含所有预期的属性，即使某些未计算
    3. 类型安全：所有值都是 float 类型，便于后续数值计算
    4. 可扩展性：如果需要添加新指标，只需修改这里的构造逻辑
    
    重要说明：
    - 如果某个 K 值没有在 k_values 中指定，对应指标将为 0.0
    - 这种设计确保了 RetrievalMetrics 对象的一致性
    - 0.0 表示"未计算"，而不是"性能为0"
    """
    
    # 确保所有必需的指标都存在，使用安全的字典访问
    result = RetrievalMetrics(
        # Image-to-Text 指标（默认为 0.0 表示未计算）
        i2t_r1=metrics_dict.get('i2t_r1', 0.0),      # I2T Recall@1
        i2t_r5=metrics_dict.get('i2t_r5', 0.0),      # I2T Recall@5
        i2t_r10=metrics_dict.get('i2t_r10', 0.0),    # I2T Recall@10
        
        # Text-to-Image 指标（默认为 0.0 表示未计算）
        t2i_r1=metrics_dict.get('t2i_r1', 0.0),      # T2I Recall@1
        t2i_r5=metrics_dict.get('t2i_r5', 0.0),      # T2I Recall@5
        t2i_r10=metrics_dict.get('t2i_r10', 0.0),    # T2I Recall@10
        
        # 平均指标（最重要的综合性能指标）
        mean_r1=metrics_dict.get('mean_r1', 0.0),    # 平均 Recall@1（最关键指标）
        mean_r5=metrics_dict.get('mean_r5', 0.0),    # 平均 Recall@5
        mean_r10=metrics_dict.get('mean_r10', 0.0),  # 平均 Recall@10
    )
    
    # ==========================================
    # 第七步：返回评估结果
    # ==========================================
    
    """
    返回的 RetrievalMetrics 对象包含了完整的评估结果，
    可以用于：
    
    1. 直接访问：result.mean_r1, result.i2t_r5 等
    2. 字典转换：result.to_dict() 用于日志记录
    3. 格式化输出：配合 print_metrics() 函数使用
    4. 性能监控：在训练循环中跟踪模型性能变化
    5. 模型比较：比较不同模型或配置的性能
    
    典型使用模式：
    - 训练时：使用 result.mean_r1 作为验证指标
    - 论文撰写：报告所有指标的完整结果
    - 调试分析：分析 I2T 和 T2I 的性能差异
    """
    
    return result


# ==========================================
# 辅助函数
# ==========================================

def compute_similarity_matrix(
    image_features: torch.Tensor,
    text_features: torch.Tensor,
    temperature: float = 1.0,
    normalize: bool = True
) -> torch.Tensor:
    """
    计算图像和文本特征之间的相似度矩阵
    
    这是一个独立的工具函数，从主要评估函数中提取出来，
    方便在需要时单独使用相似度矩阵（如可视化分析、调试等）。
    
    功能说明：
    =========
    该函数实现了 CLIP 模型中最核心的相似度计算步骤，
    支持可选的特征归一化和温度缩放。
    
    数学原理：
    =========
    1. 如果启用归一化：先对特征进行 L2 归一化
       normalized_x = x / ||x||_2
    2. 计算点积相似度：S = X @ Y^T
    3. 应用温度缩放：S_scaled = S / temperature
    
    参数详解：
    ==========
    image_features : torch.Tensor, shape [N, D]
        图像特征张量
        - N: 样本数量
        - D: 特征维度
        
    text_features : torch.Tensor, shape [N, D]
        文本特征张量，必须与图像特征具有相同的形状
        
    temperature : float, default 1.0
        温度参数，控制相似度分布的尖锐度
        - < 1.0: 分布更尖锐，差异放大
        - = 1.0: 不改变原始分布
        - > 1.0: 分布更平滑，差异缩小
        
    normalize : bool, default True
        是否对特征进行 L2 归一化
        - True: 计算余弦相似度，更稳定
        - False: 计算点积相似度，保留原始量级信息
        
    返回值：
    ========
    torch.Tensor, shape [N, N]
        相似度矩阵，其中 result[i,j] = similarity(image_i, text_j)
        
    使用示例：
    ==========
    >>> img_feat = torch.randn(10, 512)
    >>> txt_feat = torch.randn(10, 512)
    >>> sim_matrix = compute_similarity_matrix(img_feat, txt_feat)
    >>> print(f"相似度范围: [{sim_matrix.min():.3f}, {sim_matrix.max():.3f}]")
    
    应用场景：
    ==========
    - 相似度分析：分析模型学到的图像-文本语义关系
    - 可视化展示：生成相似度热力图
    - 调试诊断：检查模型是否正确对齐配对样本
    - 自定义评估：基于相似度矩阵实现其他评估指标
    """
    
    # 特征归一化（可选）
    if normalize:
        # L2 归一化：将每个特征向量的模长标准化为 1
        image_features = F.normalize(image_features, p=2, dim=1)
        text_features = F.normalize(text_features, p=2, dim=1)
    
    # 计算相似度矩阵并应用温度缩放
    # torch.matmul(A, B.t()) 等价于 A @ B^T，计算所有配对的点积
    return torch.matmul(image_features, text_features.t()) / temperature


def compute_recall_at_k(
    similarities: torch.Tensor,
    k: int,
    direction: str = 'i2t'
) -> float:
    """
    计算单个 K 值的召回率
    
    这是一个精简的召回率计算函数，专门用于计算特定 K 值和方向的召回率。
    相比主要评估函数，这个函数更加轻量，适合在需要灵活计算的场景使用。
    
    算法逻辑：
    ==========
    该函数实现了标准的 Recall@K 计算逻辑：
    1. 对相似度矩阵按指定方向进行排序
    2. 提取每个查询的 top-K 候选
    3. 检查正确答案是否在 top-K 中
    4. 计算成功的查询比例
    
    参数详解：
    ==========
    similarities : torch.Tensor, shape [N, N]
        预计算的相似度矩阵
        - similarities[i,j] = similarity(query_i, candidate_j)
        - 对角线元素是正确配对的相似度
        
    k : int
        top-K 检索的 K 值
        - 必须为正整数
        - 如果 k > N，自动返回 1.0（完美召回）
        
    direction : str, 'i2t' 或 't2i'
        检索方向
        - 'i2t': Image-to-Text，行排序，图像查询文本
        - 't2i': Text-to-Image，列排序，文本查询图像
        
    返回值：
    ========
    float
        召回率数值，范围在 [0.0, 1.0]
        - 0.0: 没有任何查询找到正确答案
        - 1.0: 所有查询都在 top-K 中找到正确答案
        
    异常处理：
    ==========
    ValueError: 当 direction 参数不是 'i2t' 或 't2i' 时抛出
    
    使用示例：
    ==========
    >>> sim = torch.randn(20, 20)  # 随机相似度矩阵
    >>> i2t_r3 = compute_recall_at_k(sim, k=3, direction='i2t')
    >>> t2i_r3 = compute_recall_at_k(sim, k=3, direction='t2i')
    >>> print(f"I2T R@3: {i2t_r3:.3f}, T2I R@3: {t2i_r3:.3f}")
    
    性能考虑：
    ==========
    - 该函数直接操作预计算的相似度矩阵，避免重复计算
    - 使用 PyTorch 原生操作，支持 GPU 加速
    - 内存占用：O(N²) 用于存储相似度矩阵
    - 时间复杂度：O(N²log(N)) 主要来自排序操作
    """
    
    n = similarities.shape[0]  # 样本数量
    
    # 边界条件：K 大于样本数时，返回完美召回率
    if k > n:
        return 1.0
    
    device = similarities.device  # 保持设备一致性
    
    if direction == 'i2t':
        # Image-to-Text 方向：对每一行（每个图像查询）进行排序
        # 找到每个图像最相似的 K 个文本
        ranked_indices = torch.argsort(similarities, dim=1, descending=True)[:, :k]
        # 构造目标：每个图像的正确文本索引（对角线位置）
        targets = torch.arange(n, device=device).unsqueeze(1)  # [N, 1]
        # 检查正确答案是否在 top-K 中
        correct = torch.any(ranked_indices == targets, dim=1)  # [N]
        
    elif direction == 't2i':
        # Text-to-Image 方向：对每一列（每个文本查询）进行排序  
        # 找到每个文本最相似的 K 个图像
        ranked_indices = torch.argsort(similarities, dim=0, descending=True)[:k, :]
        # 构造目标：每个文本的正确图像索引（对角线位置）
        targets = torch.arange(n, device=device).unsqueeze(0)  # [1, N]
        # 检查正确答案是否在 top-K 中
        correct = torch.any(ranked_indices == targets, dim=0)  # [N]
        
    else:
        # 参数验证：确保方向参数正确
        raise ValueError(
            f"Invalid direction: {direction}. Must be 'i2t' (Image-to-Text) or 't2i' (Text-to-Image)."
        )
    
    # 计算召回率：成功查询的比例
    return correct.float().mean().item()


# ==========================================
# 便捷函数
# ==========================================

def quick_eval(
    image_features: torch.Tensor,
    text_features: torch.Tensor
) -> Tuple[float, float, float]:
    """
    快速评估函数 - 返回最常用的三个综合性能指标
    
    这是一个高层次的便捷函数，为快速模型评估而设计。
    它隐藏了复杂的参数配置，直接返回最重要的平均召回率指标。
    
    适用场景：
    ==========
    - 快速原型验证：在开发过程中快速检查模型性能
    - 批量模型比较：比较多个模型的整体性能
    - 实时监控：在训练过程中快速获取关键指标
    - 简化报告：只关心综合性能而非详细指标时使用
    
    内部实现：
    ==========
    该函数内部调用完整的 compute_clip_retrieval_metrics 函数，
    使用默认参数配置（k_values=[1,5,10], temperature=1.0, normalize=True），
    然后提取三个最重要的平均指标。
    
    参数：
    ======
    image_features : torch.Tensor, shape [N, D]
        图像特征张量
    text_features : torch.Tensor, shape [N, D]  
        文本特征张量
        
    返回值：
    ========
    Tuple[float, float, float]
        返回三元组 (mean_r1, mean_r5, mean_r10)
        - mean_r1: 平均 Recall@1，最严格的综合指标
        - mean_r5: 平均 Recall@5，中等严格程度的综合指标
        - mean_r10: 平均 Recall@10，最宽松的综合指标
        
    使用示例：
    ==========
    >>> img_feat = torch.randn(50, 768)
    >>> txt_feat = torch.randn(50, 768)
    >>> r1, r5, r10 = quick_eval(img_feat, txt_feat)
    >>> print(f"模型性能: R@1={r1:.3f}, R@5={r5:.3f}, R@10={r10:.3f}")
    >>> if r1 > 0.5:
    ...     print("模型性能良好！")
    
    性能对比：
    ==========
    相比完整评估函数的优势：
    - 接口简洁：只需要特征输入，无需配置参数
    - 返回简化：只返回最关心的三个指标
    - 易于使用：适合非专业用户快速评估
    
    限制：
    - 不支持自定义参数（如温度、K值等）
    - 不提供单向指标（I2T, T2I）的详细信息
    - 不支持结果序列化和结构化访问
    """
    
    # 调用完整评估函数，使用默认参数
    metrics = compute_clip_retrieval_metrics(image_features, text_features)
    
    # 提取并返回三个最重要的综合指标
    return metrics.mean_r1, metrics.mean_r5, metrics.mean_r10


def print_metrics(metrics: RetrievalMetrics, title: str = "Retrieval Metrics"):
    """
    格式化打印评估指标 - 专业的结果展示函数
    
    这个函数提供了标准化的指标打印格式，确保评估结果的可读性和一致性。
    输出格式遵循学术论文和技术报告的惯例，便于结果分享和比较。
    
    设计理念：
    ==========
    - 层次化展示：按评估方向分组显示指标
    - 对齐格式：统一的数值格式和对齐方式
    - 重点突出：平均指标作为综合性能的总结
    - 标准精度：使用 4 位小数确保精度和可读性
    
    输出格式示例：
    ==============
    === Model Performance ===
    Image-to-Text:
      R@1:  0.4532
      R@5:  0.7234
      R@10: 0.8567
    Text-to-Image:
      R@1:  0.4123
      R@5:  0.6987
      R@10: 0.8234
    Average:
      R@1:  0.4328    <- 最重要的综合指标
      R@5:  0.7111
      R@10: 0.8401
    
    参数：
    ======
    metrics : RetrievalMetrics
        包含所有评估指标的结构化对象
        
    title : str, default "Retrieval Metrics"
        打印标题，可自定义以区分不同的评估场景
        - 例如："Validation Results", "Test Performance", "Baseline Comparison"
        
    使用示例：
    ==========
    >>> metrics = compute_clip_retrieval_metrics(img_feat, txt_feat)
    >>> print_metrics(metrics, title="模型验证结果")
    >>> print_metrics(metrics, title="Epoch 10 Performance")
    
    应用场景：
    ==========
    - 训练监控：在验证阶段打印性能指标
    - 实验记录：生成标准化的实验结果报告
    - 模型比较：并排比较不同模型的性能
    - 调试分析：快速查看模型在各个方向的表现差异
    - 演示展示：向他人展示模型性能的专业格式
    
    输出解读：
    ==========
    - Image-to-Text 指标：反映给定图像找到正确文本的能力
    - Text-to-Image 指标：反映给定文本找到正确图像的能力  
    - Average 指标：综合性能评估，通常用于模型比较和选择
    - R@1 vs R@5 vs R@10：从严格到宽松的评估标准
    
    注意事项：
    ==========
    - 该函数只负责打印，不返回任何值
    - 使用标准输出流，可以被重定向到文件
    - 数值格式固定为 4 位小数，适合大多数应用场景
    - 如需其他格式，可以直接访问 metrics 对象的属性
    """
    
    # 打印标题，使用分隔符增强可读性
    print(f"\n=== {title} ===")
    
    # Image-to-Text 检索指标
    print(f"Image-to-Text:")
    print(f"  R@1:  {metrics.i2t_r1:.4f}")    # 图像查文本的严格指标
    print(f"  R@5:  {metrics.i2t_r5:.4f}")    # 图像查文本的中等指标
    print(f"  R@10: {metrics.i2t_r10:.4f}")   # 图像查文本的宽松指标
    
    # Text-to-Image 检索指标
    print(f"Text-to-Image:")
    print(f"  R@1:  {metrics.t2i_r1:.4f}")    # 文本查图像的严格指标
    print(f"  R@5:  {metrics.t2i_r5:.4f}")    # 文本查图像的中等指标
    print(f"  R@10: {metrics.t2i_r10:.4f}")   # 文本查图像的宽松指标
    
    # 平均指标（综合性能评估）
    print(f"Average:")
    print(f"  R@1:  {metrics.mean_r1:.4f}")   # 最重要的综合指标
    print(f"  R@5:  {metrics.mean_r5:.4f}")   # 中等严格的综合指标  
    print(f"  R@10: {metrics.mean_r10:.4f}")  # 最宽松的综合指标